#!/usr/bin/env python3
"""
LLM Bridge Voice Control Node - LLM桥接语音控制节点
连接EchoKit Server并将LLM结构化命令转发给机器人控制器

改进版本特性：
- 支持 MessagePack 和 JSON 双格式解析
- 智能处理分割消息和重组
- 持久连接和自动重连机制
- 优化的命令执行时序
- 增强的错误处理和恢复
"""

import asyncio
import json
import time
import threading
from enum import Enum
from typing import Optional, Dict, Any, Callable
import re

import rclpy
from rclpy.node import Node
from rclpy.parameter import Parameter
from rclpy.qos import QoSProfile, ReliabilityPolicy, DurabilityPolicy

import websockets
from websockets.exceptions import ConnectionClosed, InvalidURI, WebSocketException

from geometry_msgs.msg import Twist
from std_msgs.msg import String
from nav_msgs.msg import Odometry

# 尝试导入 msgpack，如果没有则使用 JSON 回退
try:
    import msgpack
    MSGPACK_AVAILABLE = True
except ImportError:
    MSGPACK_AVAILABLE = False

# from .echokit_message_handler import EchoKitMessageHandler

class VoiceControlState(Enum):
    """States for the voice control system."""
    DISCONNECTED = "disconnected"
    CONNECTED = "connected"
    LISTENING = "listening"
    PROCESSING = "processing"
    EXECUTING = "executing"

class LLMBridgeVoiceControlNode(Node):
    def __init__(self):
        super().__init__('llm_bridge_voice_control_node')
        
        # 创建QoS配置
        self.qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.VOLATILE,
            depth=10
        )
        
        # 创建发布器
        self.voice_text_publisher = self.create_publisher(
            String, '/voice_text', self.qos_profile
        )
        
        self.structured_command_publisher = self.create_publisher(
            String, '/structured_voice_command', self.qos_profile
        )
        
        self.status_publisher = self.create_publisher(
            String, '/voice_control_status', self.qos_profile
        )
        
        # 添加 cmd_vel 发布器用于直接控制
        self.cmd_vel_publisher = self.create_publisher(
            Twist, '/cmd_vel', self.qos_profile
        )
        
        # WebSocket连接参数
        self.server_url = "ws://192.168.2.28:8060/ws/robot"
        self.websocket = None
        self.connected = False
        self.running = False
        
        # 状态管理
        self.current_state = VoiceControlState.DISCONNECTED
        self.command_count = 0
        self.executing_command = False
        
        # JSON 缓冲区用于处理分割的消息
        self.json_buffer = ""
        self.last_message_time = 0
        
        # 初始化EchoKit消息处理器（暂时注释掉）
        # self.echokit_handler = EchoKitMessageHandler(logger=self.get_logger())
        
        self.get_logger().info("🚀 LLM Bridge Voice Control Node 已初始化")
        self.get_logger().info(f"📡 连接到: {self.server_url}")
        self.get_logger().info("✅ 支持 MessagePack 和 JSON 双格式解析")
        self.get_logger().info("🎯 集成测试节点的成功经验")
        
        # 启动WebSocket连接
        self.start_websocket_connection()
    
    def start_websocket_connection(self):
        """启动WebSocket连接"""
        self.running = True
        
        # 在单独线程中运行异步WebSocket连接
        self.websocket_thread = threading.Thread(
            target=self.run_websocket_loop,
            daemon=True
        )
        self.websocket_thread.start()
    
    def parse_server_message(self, message):
        """解析服务器消息，支持 MessagePack 和 JSON 格式"""
        try:
            # 尝试 MessagePack 解析（服务器主要使用这种格式）
            if isinstance(message, (bytes, bytearray)):
                try:
                    if MSGPACK_AVAILABLE:
                        data = msgpack.unpackb(message, raw=False)
                        return self.handle_msgpack_event(data)
                except:
                    # 如果不是 MessagePack，尝试作为 JSON 处理
                    message = message.decode('utf-8') if isinstance(message, bytes) else message
            
            # JSON 处理逻辑
            return self.parse_json_message(message)
            
        except Exception as e:
            self.get_logger().debug(f"消息解析失败: {e}")
            return None
    
    def handle_msgpack_event(self, data):
        """处理 MessagePack 格式的服务器事件"""
        if not isinstance(data, dict):
            return None
            
        event_type = data.get('type') or list(data.keys())[0] if data else None
        
        if event_type == 'ASR':
            # ASR 结果，检查是否包含机器人命令
            text = data.get('text', '')
            self.get_logger().info(f"收到 ASR: {text}")
            return self.extract_robot_command_from_text(text)
            
        elif event_type == 'Action':
            # 直接的动作命令
            action = data.get('action', '')
            self.get_logger().info(f"收到 Action: {action}")
            return self.parse_action_command(action)
            
        elif event_type == 'StartAudio':
            # TTS 开始，检查文本中是否包含命令
            text = data.get('text', '')
            self.get_logger().info(f"收到 StartAudio: {text}")
            return self.extract_robot_command_from_text(text)
            
        return None
    
    def extract_robot_command_from_text(self, text):
        """从文本中提取机器人命令 JSON"""
        if not text:
            return None
            
        # 查找 JSON 对象
        json_patterns = [
            r'\{[^{}]*"type"[^{}]*\}',  # 简单 JSON
            r'\{(?:[^{}]|{[^{}]*})*\}',  # 嵌套 JSON
        ]
        
        for pattern in json_patterns:
            matches = re.findall(pattern, text)
            for match in matches:
                try:
                    data = json.loads(match)
                    if self.is_valid_robot_command(data):
                        return data
                except:
                    continue
        
        return None
    
    def parse_action_command(self, action):
        """解析动作命令字符串为机器人命令"""
        try:
            # 如果 action 本身就是 JSON 字符串
            if action.startswith('{') and action.endswith('}'):
                data = json.loads(action)
                if self.is_valid_robot_command(data):
                    return data
            
            # 如果是简单的动作名称，转换为标准命令
            action_map = {
                'forward': {'type': 'continuous_motion', 'linear': 0.3, 'angular': 0.0},
                'backward': {'type': 'continuous_motion', 'linear': -0.3, 'angular': 0.0},
                'left': {'type': 'continuous_motion', 'linear': 0.0, 'angular': 1.0},
                'right': {'type': 'continuous_motion', 'linear': 0.0, 'angular': -1.0},
                'stop': {'type': 'stop'},
            }
            
            if action.lower() in action_map:
                cmd = action_map[action.lower()].copy()
                cmd['original_command'] = action
                cmd['description'] = f"Execute {action}"
                return cmd
                
        except Exception as e:
            self.get_logger().debug(f"动作解析失败: {e}")
        
        return None
    
    def parse_json_message(self, message):
        """解析 JSON 消息，处理分割情况"""
        current_time = time.time()
        
        # 如果距离上次消息时间超过 1 秒，清空缓冲区
        if current_time - self.last_message_time > 1.0:
            self.json_buffer = ""
        
        self.last_message_time = current_time
        
        # 添加到缓冲区
        self.json_buffer += message
        
        # 尝试解析完整的 JSON
        return self.extract_complete_json()
    
    def extract_complete_json(self):
        """从缓冲区提取完整的 JSON 对象"""
        buffer = self.json_buffer.strip()
        
        # 查找完整的 JSON 对象
        brace_count = 0
        start_pos = -1
        
        for i, char in enumerate(buffer):
            if char == '{':
                if start_pos == -1:
                    start_pos = i
                brace_count += 1
            elif char == '}':
                brace_count -= 1
                if brace_count == 0 and start_pos != -1:
                    # 找到完整的 JSON
                    json_str = buffer[start_pos:i+1]
                    try:
                        data = json.loads(json_str)
                        if self.is_valid_robot_command(data):
                            # 清空已处理的部分
                            self.json_buffer = buffer[i+1:]
                            return data
                    except:
                        continue
        
        return None
    
    def is_valid_robot_command(self, data):
        """验证是否为有效的机器人命令"""
        if not isinstance(data, dict):
            return False
            
        cmd_type = data.get('type')
        if cmd_type not in ['timed_motion', 'continuous_motion', 'stop']:
            return False
            
        if cmd_type in ['timed_motion', 'continuous_motion']:
            if 'linear' not in data or 'angular' not in data:
                return False
            try:
                float(data['linear'])
                float(data['angular'])
            except:
                return False
                
        if cmd_type == 'timed_motion' and 'duration' not in data:
            return False
            
        return True
    
    def run_websocket_loop(self):
        """运行WebSocket事件循环"""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            loop.run_until_complete(self.websocket_handler())
        except Exception as e:
            self.get_logger().error(f"WebSocket loop error: {e}")
        finally:
            loop.close()
    
    async def websocket_handler(self):
        """简化的 WebSocket 连接处理器 - 基于测试节点的成功经验"""
        while self.running:
            try:
                self.get_logger().info("🔗 尝试连接...")
                
                async with websockets.connect(self.server_url) as ws:
                    self.websocket = ws
                    self.connected = True
                    self.get_logger().info("✅ 连接成功")
                    
                    # 注册机器人客户端
                    await ws.send(json.dumps({
                        "type": "register",
                        "client_type": "robot",
                        "robot_id": "fishbot_llm_bridge"
                    }))
                    
                    await ws.send(json.dumps({
                        "type": "robot_ready",
                        "status": "ready"
                    }))
                    
                    self.get_logger().info("📝 已注册，等待命令...")
                    self.publish_status("connected")
                    
                    # 监听消息
                    while self.running:
                        try:
                            message = await asyncio.wait_for(ws.recv(), timeout=1.0)
                            await self.handle_message(message)
                        except asyncio.TimeoutError:
                            continue
                        except websockets.exceptions.ConnectionClosed:
                            self.get_logger().warn("连接关闭")
                            break
                            
            except Exception as e:
                self.get_logger().error(f"连接错误: {e}")
                self.connected = False
                self.websocket = None
                self.publish_status("disconnected")
                await asyncio.sleep(2)
    
    async def handle_message(self, message):
        """处理消息 - 基于测试节点的成功经验"""
        try:
            self.get_logger().info(f"📨 收到消息: {message}")
            
            # 尝试解析 JSON
            if isinstance(message, str) and message.startswith('{'):
                try:
                    data = json.loads(message)
                    await self.execute_command(data)
                except json.JSONDecodeError:
                    pass
            
            # 尝试从二进制消息中提取 JSON
            elif isinstance(message, bytes):
                json_start = message.find(b'{')
                if json_start != -1:
                    json_part = message[json_start:].decode('utf-8')
                    try:
                        data = json.loads(json_part)
                        await self.execute_command(data)
                    except:
                        pass
                        
        except Exception as e:
            self.get_logger().error(f"消息处理错误: {e}")
    
    async def execute_command(self, data):
        """执行命令 - 基于测试节点的成功经验"""
        if not isinstance(data, dict):
            return
            
        cmd_type = data.get('type')
        if cmd_type not in ['timed_motion', 'continuous_motion', 'stop']:
            return
        
        # 防止重复执行
        if self.executing_command:
            self.get_logger().warn("命令执行中，跳过新命令")
            return
            
        self.executing_command = True
        self.command_count += 1
        
        try:
            original_cmd = data.get('original_command', f'Command #{self.command_count}')
            self.get_logger().info(f"🎯 执行命令 #{self.command_count}: {original_cmd}")
            
            # 发布语音文本
            self.publish_voice_text(original_cmd)
            
            # 发布结构化命令
            self.publish_structured_command(data)
            
            # 发布状态
            self.publish_status(f"executing_{cmd_type}")
            
            # 执行运动命令
            if cmd_type == 'timed_motion':
                await self.execute_timed_motion(data)
            elif cmd_type == 'continuous_motion':
                await self.execute_continuous_motion(data)
            elif cmd_type == 'stop':
                await self.execute_stop()
                
        finally:
            self.executing_command = False
            self.publish_status("ready")
    
    async def execute_timed_motion(self, data):
        """执行定时运动 - 基于测试节点的成功经验"""
        twist = Twist()
        twist.linear.x = float(data.get('linear', 0.0))
        twist.angular.z = float(data.get('angular', 0.0))
        duration = float(data.get('duration', 1.0))
        
        self.get_logger().info(f"⏱️  定时运动: linear={twist.linear.x}, angular={twist.angular.z}, duration={duration}s")
        
        # 在单独线程中执行定时运动
        def timed_motion():
            start_time = time.time()
            while time.time() - start_time < duration and self.running:
                self.cmd_vel_publisher.publish(twist)
                time.sleep(0.1)
            
            # 停止
            stop_twist = Twist()
            self.cmd_vel_publisher.publish(stop_twist)
            self.get_logger().info(f"✅ 定时运动完成 #{self.command_count}")
        
        threading.Thread(target=timed_motion, daemon=True).start()
    
    async def execute_continuous_motion(self, data):
        """执行连续运动 - 基于测试节点的成功经验"""
        twist = Twist()
        twist.linear.x = float(data.get('linear', 0.0))
        twist.angular.z = float(data.get('angular', 0.0))
        
        self.get_logger().info(f"🔄 连续运动: linear={twist.linear.x}, angular={twist.angular.z}")
        
        # 发送几次确保命令被接收
        for _ in range(3):
            self.cmd_vel_publisher.publish(twist)
            await asyncio.sleep(0.05)
    
    async def execute_stop(self):
        """执行停止 - 基于测试节点的成功经验"""
        self.get_logger().info("🛑 停止运动")
        
        stop_twist = Twist()
        for _ in range(3):
            self.cmd_vel_publisher.publish(stop_twist)
            await asyncio.sleep(0.05)
    

    
    async def handle_websocket_message(self, message):
        """处理WebSocket消息 - 使用改进的解析器"""
        try:
            # 使用改进的消息解析器
            command_data = self.parse_server_message(message)
            
            if command_data:
                await self.execute_command(command_data)
            else:
                # 如果不是机器人命令，尝试其他处理方式
                if isinstance(message, str):
                    await self.handle_text_message(message)
                elif isinstance(message, bytes):
                    await self.handle_binary_message(message)
                
        except Exception as e:
            self.get_logger().error(f"Error handling WebSocket message: {e}")
    
    async def handle_binary_message(self, message: bytes):
        """处理二进制消息（兼容旧版本）"""
        try:
            # 尝试查找 JSON 片段
            json_start = message.find(b'{')
            if json_start != -1:
                json_part = message[json_start:].decode('utf-8')
                try:
                    data = json.loads(json_part)
                    if self.is_valid_robot_command(data):
                        await self.execute_command(data)
                        return
                except:
                    pass
            
            self.get_logger().debug("无法解析二进制消息")
                
        except Exception as e:
            self.get_logger().error(f"Error handling binary message: {e}")
    
    async def handle_text_message(self, message: str):
        """处理文本消息（兼容旧版本）"""
        try:
            self.get_logger().debug(f"收到文本消息: {message}")
            
            # 尝试解析为JSON
            try:
                data = json.loads(message)
                await self.handle_json_message(data)
            except json.JSONDecodeError:
                # 如果不是JSON，当作普通文本处理
                await self.handle_plain_text_message(message)
                
        except Exception as e:
            self.get_logger().error(f"Error handling text message: {e}")
    
    async def handle_json_message(self, data: Dict[str, Any]):
        """处理JSON消息（兼容旧版本）"""
        try:
            message_type = data.get('type', 'unknown')
            
            if message_type in ['timed_motion', 'continuous_motion', 'stop']:
                # 直接执行机器人命令
                await self.execute_command(data)
                
            elif message_type == 'asr_result':
                # 语音识别结果
                text = data.get('text', '')
                self.publish_voice_text(text)
                
            elif message_type == 'llm_response':
                # LLM响应 - 可能包含结构化命令
                content = data.get('content', '')
                await self.process_llm_response(content)
                
            elif message_type == 'action':
                # 动作命令
                action = data.get('action', '')
                await self.process_action_command(action)
                
            else:
                self.get_logger().debug(f"Unknown message type: {message_type}")
                
        except Exception as e:
            self.get_logger().error(f"Error handling JSON message: {e}")
    
    async def handle_plain_text_message(self, message: str):
        """处理纯文本消息（兼容旧版本）"""
        # 检查是否是结构化命令（JSON格式）
        if message.strip().startswith('{') and message.strip().endswith('}'):
            try:
                # 尝试解析为结构化命令
                command_data = json.loads(message.strip())
                if self.is_valid_robot_command(command_data):
                    await self.execute_command(command_data)
                    return
            except json.JSONDecodeError:
                pass
        
        # 如果不是结构化命令，当作语音文本处理
        self.publish_voice_text(message)
    
    async def process_llm_response(self, content: str):
        """处理LLM响应（兼容旧版本）"""
        self.get_logger().info(f"Processing LLM response: {content}")
        
        # 检查是否是结构化命令
        if content.strip().startswith('{') and content.strip().endswith('}'):
            try:
                command_data = json.loads(content.strip())
                if self.is_valid_robot_command(command_data):
                    self.get_logger().info(f"Found structured command in LLM response: {command_data}")
                    await self.execute_command(command_data)
                    return
            except json.JSONDecodeError:
                pass
        
        # 如果不是结构化命令，发布为普通文本
        self.publish_voice_text(content)
    
    async def process_action_command(self, action: str):
        """处理动作命令（兼容旧版本）"""
        self.get_logger().info(f"Processing action command: {action}")
        
        # 使用改进的动作解析器
        command_data = self.parse_action_command(action)
        if command_data:
            await self.execute_command(command_data)
        else:
            self.publish_voice_text(action)
    
    def is_valid_structured_command(self, data: Dict[str, Any]) -> bool:
        """验证是否是有效的结构化命令（兼容旧版本）"""
        return self.is_valid_robot_command(data)
    
    def convert_action_to_structured(self, action: str) -> Optional[Dict[str, Any]]:
        """将简单动作转换为结构化命令"""
        action = action.lower().strip()
        
        # 基础命令映射
        basic_commands = {
            'forward': {'type': 'continuous_motion', 'linear': 0.3, 'angular': 0.0, 'description': 'Move forward'},
            'backward': {'type': 'continuous_motion', 'linear': -0.3, 'angular': 0.0, 'description': 'Move backward'},
            'left': {'type': 'continuous_motion', 'linear': 0.0, 'angular': 1.0, 'description': 'Turn left'},
            'right': {'type': 'continuous_motion', 'linear': 0.0, 'angular': -1.0, 'description': 'Turn right'},
            'stop': {'type': 'stop', 'description': 'Stop all motion'},
        }
        
        if action in basic_commands:
            command = basic_commands[action].copy()
            command['original_command'] = action
            return command
        
        return None
    
    def publish_voice_text(self, text: str):
        """发布语音文本"""
        msg = String()
        msg.data = text
        self.voice_text_publisher.publish(msg)
        self.get_logger().info(f"Published voice text: {text}")
    
    def publish_structured_command(self, command_data: Dict[str, Any]):
        """发布结构化命令"""
        msg = String()
        msg.data = json.dumps(command_data)
        self.structured_command_publisher.publish(msg)
        self.get_logger().info(f"Published structured command: {command_data}")
    
    def publish_status(self, status: str):
        """发布状态信息"""
        msg = String()
        msg.data = status
        self.status_publisher.publish(msg)
    
    def destroy_node(self):
        """清理资源"""
        self.running = False
        if self.websocket:
            try:
                asyncio.run(self.websocket.close())
            except:
                pass
        if hasattr(self, 'websocket_thread'):
            self.websocket_thread.join(timeout=2)
        super().destroy_node()

def main(args=None):
    print("🚀 LLM Bridge Voice Control Node")
    print("=" * 40)
    
    rclpy.init(args=args)
    
    try:
        node = LLMBridgeVoiceControlNode()
        
        print("✅ 节点启动成功")
        print("🎤 等待语音命令...")
        print("📡 监控话题:")
        print("   - /cmd_vel (运动控制)")
        print("   - /voice_text (语音文本)")
        print("   - /structured_voice_command (结构化命令)")
        print("   - /voice_control_status (状态信息)")
        print("")
        print("按 Ctrl+C 停止")
        
        rclpy.spin(node)
        
    except KeyboardInterrupt:
        print("\n🛑 停止节点")
    finally:
        if 'node' in locals():
            node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()